#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jul  5 22:31:06 2018
@project: 天池比赛-A股主板上市公司公告信息抽取
@group: MZH_314
@author: LHQ
"""
import re
import os
import sys
import decimal
from calendar import monthcalendar

import pandas as pd
from sklearn.externals import joblib

from reportIE.preprocess import ZjcReport
from reportIE.ner import CompanyAlias
from reportIE.utils import load_report
from prepare_trainingdata_zengjianchi_table import process_zjc_tables
from ltp import Ner

from train_zengjianchi_table import ColFeature

    
table_clf = joblib.load("models/zengjianchi/table_clf.m")    

ner = Ner()
columns = ["公告id", "股东全称", "股东简称", "变动截止日期", "变动价格", "变动数量", "变动后持股数", "变动后持股比例"]


def merge_tables(tbs):
    length = len(tbs)
    if length == 0:
        return pd.DataFrame()
    if length == 1:
        return tbs[0]
    elif length == 2:
        a, b = tbs
        key = list(set(a.columns) & set(b.columns))
        if len(key) == 0:
            df = pd.concat(tbs, axis=1)
        else:
            df = pd.merge(a, b, how='outer', on=key)
        return df
    else:
        tb_new = merge_tables(tbs[:2])
        return merge_tables([tb_new, *tbs[2:]])


def rebuild_table(df):
    columns = ["公告id", "股东全称", "股东简称", "变动截止日期", "变动价格", "变动数量", "变动后持股数", "变动后持股比例"]
    data = []
    for report_id, gp in df.groupby("report_id"):
        tbs = []
        for tb_index, tb in gp.groupby("table_index"):
            label = tb['label'].tolist()
            value = tb['values'].tolist()
            df_tb = pd.DataFrame.from_dict(dict(zip(label, value)), orient='columns')
            tbs.append(df_tb)
        dftb = merge_tables(tbs)
        dftb.insert(0, "公告id", report_id)
        dftb_full = dftb.reindex(columns=columns)
        data.append(dftb_full)
    dfall = pd.concat(data, ignore_index=False)
    return dfall
        

def fix_shareholder(df, html_dir):
    data = []
    for i, row in df.iterrows():
        report_id = row['公告id']
        shareholder_fullname = row['股东全称']
        shareholder_shortname = row['股东简称']
        html_path = os.path.join(html_dir, "%s.html"%report_id)
        if not os.path.exists(html_path):
            html_path = os.path.join(html_dir, "%s"%report_id)
        
        if isinstance(shareholder_fullname, float) and isinstance(shareholder_shortname, str):
            # 简称是人名则全称也是人名
            if len(shareholder_shortname) < 4:
                people_name= ner.recog_name(shareholder_shortname)
                if len(people_name) > 0:
                    row['股东全称'] = shareholder_shortname
                    row['股东简称'] = None
            else:
                _, html = load_report(html_path)
                report = ZjcReport(html, report_id)
                texts = report.flat_contents
                alias = CompanyAlias(texts, ner.recog_org)
                row['股东全称'] = alias.get_fullname_for_short(shareholder_shortname)
        ## 进一步修正全称
        shareholder_fullname = row['股东全称']
        if isinstance(shareholder_fullname, str):
            shareholder_fullname = re.sub("^\W*", "", shareholder_fullname)
            if len(shareholder_fullname) < 5:
                _, html = load_report(html_path)
                report = ZjcReport(html, report_id)
                texts = report.flat_contents
                alias = CompanyAlias(texts, ner.recog_org)
                fullname = alias.get_fullname_for_short(shareholder_fullname)
                if fullname is not None:
                    row['股东全称'] = fullname
                    row['股东简称'] = shareholder_fullname
            else:
                row['股东全称'] = shareholder_fullname            
        data.append(row)
    df_new = pd.concat(data,axis=1).T
    return df_new


def fix_date(df, pubdate_filepath):
    # 修复不完整的变动日期
    is_date = df['变动截止日期'].fillna("").map(lambda x: len(x) >=3)
    df = df[is_date]
    pubdate = ZengjianchiPubDate.from_excel(pubdate_filepath)
    def fix(row):
        report_id = str(row["公告id"])
        date = row['变动截止日期']
        row['变动截止日期'] = pubdate.fix_date(report_id, date)
        return row
    df_new = df.apply(fix, axis=1)
    return df_new


def fix_final_percent_and_quantity(df):
    # 修正持股比例
    data = []
    for id_, gp in df.groupby("公告id"):
        final_percent = gp['变动后持股比例'].unique().tolist()
        final_quantity = gp['变动后持股数'].unique().tolist()
        if len(final_percent) == 1:
            gp = gp.sort_values('变动截止日期')
            gp['变动后持股比例'] = None
            gp.loc[gp.index[-1], "变动后持股比例"] = final_percent[0]
            gp['变动后持股数'] = None
            gp.loc[gp.index[-1], "变动后持股数"] = final_quantity[0]
        data.append(gp)
    new_df = pd.concat(data)
    return new_df
        

def process_10k(df):
    ## 数值格式整理
    data = []
    for k, row in df.iterrows():
        field = row['field']
        label = row['label']
        values= row['values']
        if label in ('变动数量', '变动后持股数'):
            new_values = []
            u = decimal.Decimal(10000) if "万" in field else decimal.Decimal(1)
            for v in values:
                v = re.sub("[,\s]", "", v)
                try:
                    v_new = decimal.Decimal(v)
                except decimal.InvalidOperation:
                    v_new = v
                else:
                    v_new *= u
                new_values.append(v_new)
        elif label in ("变动价格", "变动后持股比例"):
            new_values = []
            for v in values:
                v = re.sub("[,\s]", "", v)
                try:
                    v_new = decimal.Decimal(v)
                except decimal.InvalidOperation:
                    v_new = v
                new_values.append(v_new)
        else:
            new_values = values
        row['values'] = new_values
        data.append(row)
    new_df = pd.concat(data, axis=1).T    
    return new_df



class ZengjianchiPubDate:
    """
    增减持公告的发布日期
    数据主要应用于“股东增减持”类型公告的抽取，对于“变动截止日期”字段，存在少量公告中只公布了月份，未公布具体的日期。对这种情况的处理标准为：
    (1).如果该月份在公告发布月份的前面，变动截止日期为该月份最后1个交易日；
    (2).如果该月份是公告发布的月份，变动截止日期为公告发布日期
    """
    def __init__(self, d):
        self.d = {str(k):v for k, v in d.items()}
        
    def fix_date(self, report_id:str, date:str):
        pub_date = self.d.get(report_id)
        if not isinstance(date, str) or not isinstance(pub_date, str):
            return date
        if re.search("^\d{1,2}-\d{1,2}$", date):
            y = pub_date.split('-')[0]
            date = y + "-" + date
        elif re.search("^\d{4}-\d{1,2}$", date):
            py, pm, _ = pub_date.split("-")
            y, m = date.split("-")
            if y < py or m < pm:
                # 月底最后一个周五
                d = [w[4] for w in monthcalendar(int(y), int(m)) if w[4] != 0][-1]
                date = "%s-%s-%s" % (y, m, d)
            elif m == pm:
                date = pub_date
            else:
                pass
        else:
            pass
        return date
    
    @classmethod
    def from_excel(cls, excel_path):
        df = pd.read_excel(excel_path, index_col=1)
        d = df['公告日期'].map(str).map(lambda x: x.split()[0]).to_dict()
        pubdate = cls(d)
        return pubdate


if __name__ == "__main__":
    if len(sys.argv) == 4:
        html_dir = sys.argv[1]
        pubdate_path = sys.argv[2]
        save_path = sys.argv[3]
    else:
        ## train data
        #html_dir = os.path.abspath("../data/round1_train_20180518/增减持/html")
        #pubdate_path = "../data/[new] FDDC_announcements_round1_public_time_20180629/增减持公告日期_ 训练数据.xlsx"
        #save_path = os.path.abspath("../data/tmp/zengjianchi_table.txt")
       
        # test_data
        html_dir = os.path.abspath("../data/FDDC_announcements_round1_test_b_20180708/增减持/html")
        pubdate_path = os.path.abspath("../data/FDDC_announcements_round1_test_b_public_20180708.xlsx")
        save_path = os.path.abspath("../submit/zengjianchi_table.txt")
    
    #================================================
    save_dir = os.path.dirname(save_path)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    df = process_zjc_tables(html_dir)
    
    fields = df['field'].map(str).tolist()
    values = df['values'].tolist()
    field_index = df['field_index'].values.reshape((-1, 1))
    
    col_feat = ColFeature.from_modelfile()
    
#    features = col_feat.build_feature(fields, values, field_index)
    features = col_feat.build_feature(fields, values)
    df['label'] =  table_clf.predict(features)
    
    df_target = df.query("label != '其他'")
    df_processed_values = process_10k(df_target)
    df_full = rebuild_table(df_processed_values)
    df_col_0_2 = fix_shareholder(df_full, html_dir)
    df_col_0_3 = fix_date(df_col_0_2, pubdate_path)
    df_col_0_end = fix_final_percent_and_quantity(df_col_0_3)
    df_result = df_col_0_end.dropna(subset=["股东全称"])

    df_result.to_csv(save_path, index=False, sep='\t', line_terminator='\r\n')
